﻿CREATE OR REPLACE FUNCTION public.cluster_reads(input_query text, output_table text) RETURNS VOID AS $$

# -----------------------------[ cluster_reads.sql ]---------------------------
"""
Groups overalapping reads, with the same orientation, into clusters. I.e,
generates contigs.

USAGE:
SELECT cluster_reads(
  'SELECT chr, strand, read_start, read_end FROM my_table',
  'my_output_table'
  );

INPUT:
A table where the first four columns are: 
    1. chromosome, 
    2. strand orientation, 
    3. read start, 
    4. read end

Column names don't matter. Column order does. Change indexes below as 
appropriate if a different column position is required.
Columns other than the above are ignored.

See below for overlapping criterion

OUTPUT:
A table like the input one with a column of tag clusters added.
Note: Currently, column names in input table are changed to lowercase in the 
output

"""

import psycopg2

#---------------------------------[ User's input ]-----------------------------

## Column indexes in Python base. First column has index 0 Change as appropriate
chr_column= 0
strand_column = 1
read_start_column= 2
read_end_column= 3

## Cluster reads overalpping by at least...
overlap_length = -10  ## +ve : Cluster reads overlapping by at least n bases
                    ## -ve : Cluster reads with a gap of no more than -n bases
                    ## 0 : Cluster reads where, at least, the end of one read
                    ##     correspond to the start of the next. I.e. no gap,
                    ##     no overlap. Adjacent reads.
                       
# ----------------------------[ Connect to database ]--------------------------

conn = psycopg2.connect("""dbname='Vitelleschi' user='dberaldi' \
    host='localhost' password='mypassword'""")

cur_fetch = conn.cursor()
cur_modify= conn.cursor()

# --------------------------[ Order input table ]------------------------------

cur_modify.execute(""" CREATE TABLE tmp_cluster_reads_input AS (
  %(input_query)s) """ %{'input_query':input_query})

## Extract names from first 4 columns of the input query
cur_fetch.execute("""
  SELECT array_to_string(ARRAY(SELECT column_name::text 
  FROM information_schema.columns 
  WHERE table_name like 'tmp_cluster_reads_input' AND
    ordinal_position IN (%(chr_column)i, 
                         %(strand_column)i, 
                         %(read_start_column)i, 
                         %(read_end_column)i)), ',') 
  """ %{'chr_column':chr_column + 1 , 'strand_column':strand_column + 1, 
        'read_start_column':read_start_column + 1, 'read_end_column':read_end_column + 1})

## Format list of column names
col_names= str(list(cur_fetch.fetchone())).strip('[]')
col_names= col_names.replace("'", '"')
col_names= col_names.replace(',', '", "')

## Compile and run the create table query with correct order.
sqlstm= """ CREATE TABLE tmp_cluster_reads_ordered AS (
  SELECT * FROM tmp_cluster_reads_input ORDER BY """ + col_names + ')'
  
cur_fetch.execute(sqlstm)

cur_modify.execute(""" DROP TABLE tmp_cluster_reads_input """)

#-----------------------[ Output table ]---------------------------------------

cur_modify.execute("""CREATE TABLE %(output_table)s AS (
  SELECT *, 'cluster' AS tag_cluster 
  FROM tmp_cluster_reads_ordered
  WHERE 1=2) """ %{'output_table':output_table})

cur_modify.execute("""
 alter table %(output_table)s alter column tag_cluster TYPE varchar
  """ %{'output_table':output_table})

# ------------------[ Start iterating over input rows ]--------------------

## Start cursor to go thorugh ordered input table
cur_fetch.execute(""" SELECT * FROM tmp_cluster_reads_ordered """)

line= cur_fetch.fetchone()
line= list(line)
line= [str(x) for x in line]

previous_chr= line[chr_column] + line[strand_column]
previous_end= line[read_end_column]

## Cluster name composed by: chromosome, strand, leftmost position of the cluster.
cluster_name= line[chr_column] + '_' + line[strand_column] + '_' + line[read_start_column]

line.append(cluster_name)

line_cluster= str(line).strip('[]')

cur_modify.execute("""
  INSERT INTO %(output_table)s 
  VALUES (%(line_cluster)s) 
  """ %{'output_table':output_table, 'line_cluster':line_cluster})

while 1<2:
    line= cur_fetch.fetchone()

    if line == None:
        break

    line= list(line)
    line= [str(x) for x in line]
    
    read_end= line[read_end_column]
    read_start= line[read_start_column]

    overlap= int(previous_end) - int(read_start)

    ## Update cluster_name when a gap between reads is found or when the read orientation changes    
    if  overlap >= (overlap_length - 1) and line[chr_column] + line[strand_column] == previous_chr:
        pass
    else:
        cluster_name = line[chr_column] + '_' + line[strand_column] + '_' + line[read_start_column]

    ## This 'if' makes small reads, contained into large ones, are clustered together
    if int(previous_end) < int(line[read_end_column]) or line[chr_column] != previous_chr:
        previous_end= line[read_end_column]
        
    line.append(cluster_name)

    line_cluster= str(line).strip('[]')

    cur_modify.execute("""
      INSERT INTO %(output_table)s 
      VALUES (%(line_cluster)s) 
      """ %{'output_table':output_table, 'line_cluster':line_cluster})

    previous_chr = line[chr_column] + line[strand_column]
      
#------------------------------------------------------------------------------

cur_modify.execute(""" DROP TABLE tmp_cluster_reads_ordered """)

conn.commit()

$$ LANGUAGE 'plpythonu';